import json
from datetime import datetime, timedelta

import numpy as np
import pandas as pd

from air_web.config.config import config
from air_web.data_platform import init_db
from air_web.web_flask.dal.base_dal import EsBaseDal

sql_engine = init_db()
es_dal = EsBaseDal(config["ES_HOST"])
default_user_file = "./default_user.csv"
cons_no_field = "cons_no"


def save_json(res_list, json_name):
    with open(json_name, "w") as f:
        json.dump(res_list, f)


def proc_city_res(alert_df, today, baseline_day, baseline_coef):
    sql = f"""
              SELECT a.*, org_name as cityOrgName, b.p_kt_baseline FROM
            ((
                 SELECT
                 org_no as cityOrgNo,
                 max_p_kt p_kt_today
            FROM
                 city_report
            WHERE
                    data_time =  '{today}' and org_no in (select org_no from real_org_no where org_level=1) and type_id=0
            group by cityOrgNo
            ) a
            left join (SELECT
                 org_no as cityOrgNo,
                 max_p_kt p_kt_baseline
            FROM
                 city_report
            WHERE
                    data_time = '{baseline_day}' and org_no in (select org_no from real_org_no where org_level=1) and type_id=0
            group by cityOrgNo) b on a.cityOrgNo = b.cityOrgNo)
            left join real_org_no r on a.cityOrgNo=r.org_no
           """

    city_res_df = sql_engine.query(sql)

    if alert_df.empty:
        city_res_df["cycleCoad"] = 0.0
        city_res_df["cycleLoadType"] = 2
    else:
        city_res_df["p_kt_baseline"] = (
            city_res_df["p_kt_baseline"] * baseline_coef
        )  # 因为基线太小，客户要求乘个系数
        city_res_df["cycleCoad"] = (
            city_res_df["p_kt_baseline"] - city_res_df["p_kt_today"]
        )
        city_res_df["cycleLoadType"] = 1

        on5_list = alert_df["on5"].tolist()
        city_res_df.loc[
            ~city_res_df["cityOrgNo"].isin(on5_list), "cycleCoad"
        ] = 0.0
        city_res_df.loc[
            ~city_res_df["cityOrgNo"].isin(on5_list), "cycleLoadType"
        ] = 2

    city_res_df = city_res_df.drop(["p_kt_today", "p_kt_baseline"], axis=1)
    city_res_df["cityOrgName"] = (
        city_res_df["cityOrgName"]
        .str.replace("国网", "")
        .str.replace("供电公司", "")
        .str.replace("新区", "")
    )
    city_res_df["dayTime"] = today

    res_list = city_res_df.to_dict("record")
    return res_list


def get_city_cons_kt(on5, cid_list, data_date):
    # 如果展示数据是今天，那么end_time取现在(因为用户计算结果中存了全天的预测点)
    if datetime.strptime(data_date, "%Y-%m-%d") == datetime.now().date():
        end_date = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    else:
        end_date = (
            datetime.strptime(data_date, "%Y-%m-%d") + timedelta(days=1)
        ).strftime("%Y-%m-%d")

    rules = [
        ("data_time", "query", ">=", data_date),
        ("data_time", "query", "<", end_date),
        (cons_no_field, "query", "in", cid_list),
    ]
    is_cal = "1" if config["c_cons_filter_is_cal"] else "*"
    index_name = (
        config["POWERLOAD"]
        .replace("on5", str(on5))
        .replace("on7", "*")
        .replace("cal01", is_cal)
    )
    cons_df = es_dal.query_dataframe(
        rules,
        index_name,
        doc_time_field="data_time",
        source=[cons_no_field, "p_kt", "data_time"],
    )
    if cons_df.empty:
        print(f"查询数据为空,data_time:{data_date},index:{index_name},on5:{on5}...")
        return cons_df
    cons_df = cons_df.sort_values(
        [cons_no_field, "data_time"], ascending=False
    )
    cons_df["data_time"] = cons_df["data_time"].dt.strftime("%H:%M:%S")

    return cons_df


def get_curve_data(subst_cons_df, today, baseline_day):
    on5_dict = (
        subst_cons_df[["cityOrgNo", "consNo"]]
        .groupby("cityOrgNo")["consNo"]
        .agg(list)
        .to_dict()
    )

    subst_curve_df = pd.DataFrame()
    for on5, cid_list in on5_dict.items():
        today_df = get_city_cons_kt(on5, cid_list, today)
        baseline_df = get_city_cons_kt(on5, cid_list, baseline_day)

        curve_df = pd.merge(
            today_df,
            baseline_df,
            on=["cons_no", "data_time"],
            how="right",
            suffixes=("_today", "_baseline"),
        )
        subst_curve_df = pd.concat([subst_curve_df, curve_df])
    print(
        f"""查询用户数量:{len(subst_curve_df["cons_no"].drop_duplicates())},data_time:{today}"""
    )

    subst_curve_df = subst_curve_df.rename(
        columns={
            "data_time": "recordTime",
            "p_kt_today": "realLoad",
            "p_kt_baseline": "baseLoad",
        }
    )
    subst_curve_df["recordTime"] = f"{today} " + subst_curve_df["recordTime"]
    subst_curve_df=subst_curve_df.rename(columns={cons_no_field: "consNo"})
    subst_curve_df = subst_curve_df.sort_values(["consNo", "recordTime"])

    return subst_curve_df


def get_default_cons_data(day_time, baseline_day):
    # cons_no_list = ["9519000117", "1046469383", "0730975382", "1147913983", "1109156272"]
    cons_df = pd.read_csv(default_user_file, header=None, dtype=str)
    cons_no_list = cons_df[0].drop_duplicates().tolist()

    sql = f"""select cons_no as consNo,
                     cons_name as consName,
                     on5 as cityOrgNo,
                     shi as cityOrgName,
                     type_code as industry,
                     '{day_time}' as dayTime
              from c_cons
              where cons_no in ({','.join(cons_no_list)})
           """
    subst_cons_df = sql_engine.query(sql)

    res_subst_curve_df = get_curve_data(subst_cons_df, day_time, baseline_day)

    # 过滤基线负荷为0的点
    res_subst_curve_df = res_subst_curve_df.loc[res_subst_curve_df['baseLoad'] != 0]

    # 求调节率
    res_subst_curve_df["adjustRate"] = (
        res_subst_curve_df["baseLoad"] - res_subst_curve_df["realLoad"]
    ) / res_subst_curve_df["baseLoad"]
    res_subst_curve_df = res_subst_curve_df.replace({np.nan: 0})
    # avg_subst_curve_df = res_subst_curve_df.loc[
    #     res_subst_curve_df["adjustRate"] != -np.inf
    # ]
    # 求调节率平均值，且只取调节率小于0的用户
    mean_rate_df = (
        res_subst_curve_df.groupby("consNo")
        .agg({"adjustRate": "mean"})
        .reset_index()
    )
    mean_rate_df = mean_rate_df.loc[mean_rate_df["adjustRate"] < 0]
    subst_cons_df = pd.merge(
        subst_cons_df,
        mean_rate_df,
        on="consNo",
        how="right",
    )
    subst_cons_df.loc[subst_cons_df["adjustRate"] < -1, "adjustRate"] = -1
    res_subst_curve_df = res_subst_curve_df.loc[
        res_subst_curve_df["consNo"].isin(mean_rate_df["consNo"])
    ]
    res_subst_curve_df = res_subst_curve_df.drop("adjustRate", axis=1)

    subst_cons_df['adjustRate'] = None  # 24.6.20 南瑞要求
    return subst_cons_df, res_subst_curve_df


def get_subst_cons_df(day_time):
    sql = f"""select sc.cons_no as consNo,
                    cons_name as consName,
                    sc.on5 as cityOrgNo ,
                    r.org_name as cityOrgName,
                    rate as adjustRate,
                    cc.industry,
                    '{day_time}' as dayTime
             from substandard_cons sc
             left join
             (select cons_no, type_code as industry
              from c_cons
             )cc on sc.cons_no=cc.cons_no
             left join real_org_no r on sc.on5=r.org_no
             where compute_date='{day_time}' and rate is not null
             and compute_date-control_time < 3
             order by rate limit 50
             """
    subst_cons_df = sql_engine.query(sql)
    return subst_cons_df


def get_city_alert(day_time):
    sql = """select c.org_no as on5 from city_approval_info  c
             right join (select org_no from real_org_no where org_level=1 )r on r.org_no=c.org_no 
             and end_time >= '{today}'
          """.format(today=day_time)
    alert_df = sql_engine.query(sql)
    return alert_df


def get_no_excute_cons(day_time, baseline_day):
    baseline_day = (datetime.strptime(day_time, '%Y-%m-%d') - timedelta(days=1)).strftime('%Y-%m-%d')
    subst_cons_df = get_subst_cons_df(day_time)
    if subst_cons_df.empty:
        print(f"未查询到调节未达标用户,返回默认用户,data_time:{day_time} ...")
        subst_cons_df, res_subst_curve_df = get_default_cons_data(
            day_time, baseline_day
        )
    else:
        res_subst_curve_df = get_curve_data(
            subst_cons_df, day_time, baseline_day
        )

    res_subst_curve_df = res_subst_curve_df.replace({np.nan: None})
    res_subst_curve_list = res_subst_curve_df.to_dict("record")

    res_subst_cons_list = subst_cons_df.to_dict("record")
    print(
        f"""no_excute_cons_count:{len(res_subst_cons_list)},data_time:{day_time}"""
    )
    return res_subst_cons_list, res_subst_curve_list


def get_city_cycle_coad(day_time, baseline_day, baseline_coef):
    alert_df = get_city_alert(day_time)
    if alert_df.empty:
        print("未查询到今日预警地市,data_time:{day_time} ...")
    city_res_list = proc_city_res(
        alert_df, day_time, baseline_day, baseline_coef
    )
    return city_res_list


def test_main(baseline_coef=1):
    get_no_excute_cons(day_time, baseline_day)
    get_city_cycle_coad(day_time, baseline_day, baseline_coef)


if __name__ == "__main__":
    day_time = "2022-12-14"
    baseline_day = "2022-12-15"
    test_main()
